package team.xiaobaiya.gptchatrobot.handler.chat;

import com.squareup.okhttp.Response;
import com.squareup.okhttp.ResponseBody;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import team.xiaobaiya.gptchatrobot.exception.ChatException;
import team.xiaobaiya.gptchatrobot.gpt.GptClient;
import team.xiaobaiya.gptchatrobot.gpt.request.ChatRequestBody;
import team.xiaobaiya.gptchatrobot.gpt.request.MessageItem;
import team.xiaobaiya.gptchatrobot.rights.RightsManager;
import team.xiaobaiya.gptchatrobot.rights.UserKeyCheckResult;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * 对话封装，客户端每次对话请求都会创建这个类的对象
 *
 * @ClassName ConverseHandleWrapper
 * @Author Feng Yingnan
 * @Email fengyn1211@163.com
 * @Date 2023/4/24 9:24
 */
@Slf4j
public class ConverseHandleWrapper{
    public static GptClient gptClient;
    public static RightsManager rightsManager;
    private static final ExecutorService executorService = Executors.newFixedThreadPool(10);
    private final static Pattern contentPattern = Pattern.compile("\"content\":\"(.*?)\"}");
    private final static Integer MSG_LIST_MAX = 15;
    private final static String EVENT_DATA = "d";
    private final static String EVENT_ERROR = "e";

    // 用于数据传输的 SseEmitter
    private final SseEmitter emitter = new SseEmitter(0L);
    // 用户唯一标识
    private String userKey;
    // 对话上下文
    private List<MessageItem> messageItemList;


    /**
     * 向客户端发送数据
     * @param event 事件类型
     * @param data 数据
     */
    private boolean sendData2Client(String event, String data) {
        try {
            emitter.send(SseEmitter.event().name(event).data("{" + data + "}"));
            return true;
        } catch (IOException e) {
            log.error("向客户端发送消息时出现异常");
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 发送错误消息并结束
     * @param errorMsg 错误消息
     * @return SseEmitter
     */
    public SseEmitter sendErrorAndEnd(String errorMsg) {
        sendData2Client(EVENT_ERROR, errorMsg);
        emitter.complete();
        return emitter;
    }


    /**
     * 对话上下文检查
     * @return 是否通过
     */
    private boolean messageListCheck() {
        if (messageItemList.size() > MSG_LIST_MAX) {
            sendData2Client(EVENT_ERROR, "上下文太多了，我脑子炸了");
            return false;
        }
        return true;
    }

    /**
     * 对话处理
     * @return SseEmitter
     */
    public SseEmitter handle() {
        if (!messageListCheck() || !authenticate()) {
            return emitter;
        }
        rightsManager.lockUserKey(userKey);
        doConverse();
        return emitter;
    }


    /**
     * 鉴权
     * @return 是否通过
     */
    public boolean authenticate() {
        UserKeyCheckResult userKeyCheckResult = rightsManager.userKeyCheck(userKey);
        if (userKeyCheckResult.notOk()) {
            sendData2Client(EVENT_ERROR, userKeyCheckResult.msg);
            emitter.complete();
            return false;
        }
        return true;
    }


    /**
     * 对话，异步的，在新的线程的
     */
    public SseEmitter doConverse() {
        executorService.execute(this::run);
        return emitter;
    }

    private void run() {
        ChatRequestBody chatRequestBody = new ChatRequestBody();
        chatRequestBody.setMessages(messageItemList);
        Response chatResponse;
        try {
            chatResponse = gptClient.chat(chatRequestBody);
        } catch (ChatException e) {
            sendData2Client(EVENT_ERROR, "我累垮了");
            emitter.complete();
            e.printStackTrace();
            return;
        }
        try (ResponseBody responseBody = chatResponse.body();
             InputStream inputStream = responseBody.byteStream();
             BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream))) {
            String line;
            while ((line = bufferedReader.readLine()) != null) {
                if (StringUtils.hasLength(line)) {
                    Matcher matcher = contentPattern.matcher(line);
                    if (matcher.find()) {
                        String content = matcher.group(1);
                        if (!sendData2Client(EVENT_DATA, content)) {
                            break;
                        }
                    }
                }
            }
        } catch (IOException e) {
            log.error("ResponseBody读取错误");
            e.printStackTrace();
        } finally {
            emitter.complete();
            rightsManager.decrementUsage(userKey);
            rightsManager.unlockUserKey(userKey);
        }
    }

    public ConverseHandleWrapper() {}
    public ConverseHandleWrapper(String userKey) {
        this.userKey = userKey;
    }
    public ConverseHandleWrapper(String userKey, List<MessageItem> messageItemList) {
        this.userKey = userKey;
        this.messageItemList = messageItemList;
    }

    public SseEmitter getEmitter() {
        return emitter;
    }

    public void messageItemList(List<MessageItem> messageItemList) {
        this.messageItemList = messageItemList;
    }
    public void userKey(String userKey) {
        this.userKey = userKey;
    }
}
